source activate ke2torch23cu121

model=llama2-7b-chat
method=ICE

DATE=$(date +"%Y-%m-%d")
if [ ! -d "log/${DATE}" ]; then
    mkdir -p log/${DATE}
    mkdir -p log/${DATE}/${method}-js
    mkdir -p log/${DATE}/${method}-kl
fi

gpu_id=2

# ICE KnowEdit datasets
declare -A datasets=(
    ["zsre"]="./dataset/ICE/zsre.json"
    # ["recent"]="./dataset/ICE/wikidata_recent.json"
    ["counterfact"]="./dataset/ICE/wikidata_counterfact.json"
    # ["wikibio"]="./dataset/ICE/wikibio.json"
)

for method_type in kl js; do
    for datatype in "${!datasets[@]}"; do
        data_dir="${datasets[$datatype]}"
        echo "${DATE}/${method}-${method_type}/${model}-${datatype}-${gpu_id}"
        CUDA_VISIBLE_DEVICES=${gpu_id} python examples/run_knowedit_llama2.py \
            --editing_method=${method} \
            --hparams_dir=./hparams/${method}/${model}-${method_type}.yaml \
            --data_dir=${data_dir} \
            --datatype=${datatype} \
            --metrics_save_dir=./results/${model}/${method}-${method_type}-${DATE} \
            --pre_file=./pre_edit/${model}_${datatype}_pre_edit.json \
            > log/${DATE}/${method}-${method_type}/${model}-${datatype}-${gpu_id}.log 2>&1 &
        wait
    done
done